#include "UniASTLexer.h"
#include "UniASTParser.h"
#include "UniASTVisitor.h"
#include "antlr4-runtime.h"
#include <filesystem>
#include <fstream>
#include <iostream>

using namespace antlr4;
using namespace std;

int main(int argc, char **argv) {
  std::string output_dir(".");
  std::string input_file(argv[1]);
  for (int i = 2; i < argc; i++) {
    std::string key(argv[i]);
    if (key == "-o") {
      i++;
      if (i < argc) {
        output_dir = std::string(argv[i]);
      } else {
        std::cerr << "Please enter output dir." << std::endl;
        return 1;
      }
      i++;
    } else {
      std::cerr << "unknown key: " << key << std::endl;
      return 1;
    }
  }

  std::filesystem::path inputfilePath(input_file);
  if (!std::filesystem::exists(inputfilePath)) {
    std::cerr << "Can not find input file: " << input_file << std::endl;
    return 1;
  }
  std::filesystem::path outputDirPath = std::filesystem::path(output_dir);
  if (!std::filesystem::exists(outputDirPath)) {
    std::cerr << "Can not find output directory: " << output_dir << std::endl;
    return 1;
  }
  if (!std::filesystem::is_directory(outputDirPath)) {
    std::cerr << output_dir << " is not a directory." << std::endl;
    return 1;
  }
  std::ifstream stream;
  stream.open(input_file);
  ANTLRInputStream input(stream);
  UniAST::UniASTLexer lexer(&input);
  CommonTokenStream tokens(&lexer);
  UniAST::UniASTParser parser(&tokens);
  auto root = parser.rules();

  auto grammarFileName = inputfilePath.filename().stem().string();
  auto lexerFilePath =
      outputDirPath / (std::filesystem::path(grammarFileName + "Lexer.g4"));
  auto parserFilePath =
      outputDirPath / (std::filesystem::path(grammarFileName + "Parser.g4"));
  auto listenerHeaderFilePath =
      outputDirPath / (std::filesystem::path(grammarFileName + "Listener.h"));
  auto listenerCppFilePath =
      outputDirPath / (std::filesystem::path(grammarFileName + "Listener.cpp"));
  auto cmakeFilePath = outputDirPath / std::filesystem::path("CMakeLists.txt");
  auto driverFilePath = outputDirPath / std::filesystem::path("Driver.cpp");
  std::ofstream lexerFileStream(lexerFilePath);
  std::ofstream parserFileStream(parserFilePath);
  std::ofstream listenerHeaderFileStream(listenerHeaderFilePath);
  std::ofstream listenerCppFileStream(listenerCppFilePath);
  std::ofstream cmakeFileStream(cmakeFilePath);
  std::ofstream driverStream(driverFilePath);
  UniAST::Visitor v(grammarFileName, lexerFileStream, parserFileStream,
                    listenerHeaderFileStream, listenerCppFileStream,
                    cmakeFileStream, driverStream);
  v.visit(root);
  return 0;
}